# Copyright (C) 2025 Kumo inc.
# Author: Jeff.li lijippy@163.com
# All rights reserved.
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published
# by the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program.  If not, see <https:#www.gnu.org/licenses/>.
#

"""Tests for test randomization."""

import random
import subprocess

from turbo import flags
from turbo.testing import _bazelize_command
from turbo.testing import turbotest
from turbo.testing import parameterized
from turbo.testing.tests import turbotest_env

FLAGS = flags.FLAGS


class TestOrderRandomizationTest(parameterized.TestCase):
  """Integration tests: Runs a py_test binary with randomization.

  This is done by setting flags and environment variables.
  """

  def setUp(self):
    super().setUp()
    self._test_name = 'turbo/testing/tests/turbotest_randomization_testcase'

  def _run_test(self, extra_argv, extra_env):
    """Runs the py_test binary in a subprocess, with the given args or env.

    Args:
      extra_argv: extra args to pass to the test
      extra_env: extra env vars to set when running the test

    Returns:
      (stdout, test_cases, exit_code) tuple of (str, list of strs, int).
    """
    env = turbotest_env.inherited_env()
    # If *this* test is being run with this flag, we don't want to
    # automatically set it for all tests we run.
    env.pop('TEST_RANDOMIZE_ORDERING_SEED', '')
    if extra_env is not None:
      env.update(extra_env)

    command = (
        [_bazelize_command.get_executable_path(self._test_name)] + extra_argv)
    proc = subprocess.Popen(
        args=command,
        env=env,
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
        universal_newlines=True)

    stdout, _ = proc.communicate()

    test_lines = [l for l in stdout.splitlines() if l.startswith('class ')]
    return stdout, test_lines, proc.wait()

  def test_no_args(self):
    output, tests, exit_code = self._run_test([], None)
    self.assertEqual(0, exit_code, msg='command output: ' + output)
    self.assertNotIn('Randomizing test order with seed:', output)
    cases = ['class A test ' + t for t in ('A', 'B', 'C')]
    self.assertEqual(cases, tests)

  @parameterized.parameters(
      {
          'argv': ['--test_randomize_ordering_seed=random'],
          'env': None,
      },
      {
          'argv': [],
          'env': {
              'TEST_RANDOMIZE_ORDERING_SEED': 'random',
          },
      },)
  def test_simple_randomization(self, argv, env):
    output, tests, exit_code = self._run_test(argv, env)
    self.assertEqual(0, exit_code, msg='command output: ' + output)
    self.assertIn('Randomizing test order with seed: ', output)
    cases = ['class A test ' + t for t in ('A', 'B', 'C')]
    # This may come back in any order; we just know it'll be the same
    # set of elements.
    self.assertSameElements(cases, tests)

  @parameterized.parameters(
      {
          'argv': ['--test_randomize_ordering_seed=1'],
          'env': None,
      },
      {
          'argv': [],
          'env': {
              'TEST_RANDOMIZE_ORDERING_SEED': '1'
          },
      },
      {
          'argv': [],
          'env': {
              'LATE_SET_TEST_RANDOMIZE_ORDERING_SEED': '1'
          },
      },
  )
  def test_fixed_seed(self, argv, env):
    output, tests, exit_code = self._run_test(argv, env)
    self.assertEqual(0, exit_code, msg='command output: ' + output)
    self.assertIn('Randomizing test order with seed: 1', output)
    # Even though we know the seed, we need to shuffle the tests here, since
    # this behaves differently in Python2 vs Python3.
    shuffled_cases = ['A', 'B', 'C']
    random.Random(1).shuffle(shuffled_cases)
    cases = ['class A test ' + t for t in shuffled_cases]
    # We know what order this will come back for the random seed we've
    # specified.
    self.assertEqual(cases, tests)

  @parameterized.parameters(
      {
          'argv': ['--test_randomize_ordering_seed=0'],
          'env': {
              'TEST_RANDOMIZE_ORDERING_SEED': 'random'
          },
      },
      {
          'argv': [],
          'env': {
              'TEST_RANDOMIZE_ORDERING_SEED': '0'
          },
      },)
  def test_disabling_randomization(self, argv, env):
    output, tests, exit_code = self._run_test(argv, env)
    self.assertEqual(0, exit_code, msg='command output: ' + output)
    self.assertNotIn('Randomizing test order with seed:', output)
    cases = ['class A test ' + t for t in ('A', 'B', 'C')]
    self.assertEqual(cases, tests)


if __name__ == '__main__':
  turbotest.main()
